import cv2
import os

import numpy as np

# 获取训练数据集目录的路径
data_path = 'dataset'

# 收集数据集中的人脸图像和标签
faces = []
labels = []
label_dict = {}

# 遍历数据集文件夹，处理每个人的图像
for root, dirs, files in os.walk(data_path):
    for file in files:
        if file.endswith('.jpg') or file.endswith('.png'):
            # 提取人脸标签
            label = os.path.basename(root).replace(" ", "-").lower()
            if label not in label_dict:
                label_dict[label] = len(label_dict) + 1

            # 加载图像并进行灰度转换
            image_path = os.path.join(root, file)
            image = cv2.imread(image_path)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

            # 使用人脸检测器检测人脸
            face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
            faces_rect = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

            # 提取人脸并添加到训练集
            for (x, y, w, h) in faces_rect:
                face = gray[y:y+h, x:x+w]
                faces.append(face)
                labels.append(label_dict[label])

# 创建LBPH人脸识别器并训练模型
recognizer = cv2.face.LBPHFaceRecognizer_create()
recognizer.train(faces, np.array(labels))

# 保存训练好的模型
recognizer.save('face_model.xml')